import os
import importlib.metadata
from pathlib import Path


def get_package_size_improved(package_name):
    try:
        # 使用新的 importlib.metadata 模块
        dist = importlib.metadata.distribution(package_name)
        package_path = Path(dist._path).parent

        # OpenCV 包的实际目录名通常是 'cv2'
        cv2_dir = package_path / 'cv2'

        if cv2_dir.exists():
            total_size = 0
            for filepath in cv2_dir.rglob('*'):
                if filepath.is_file():
                    total_size += filepath.stat().st_size

            # 转换为 MB 显示
            size_mb = total_size / (1024 * 1024)
            return f"{size_mb:.2f} MB"
        else:
            return f"Package directory not found: {cv2_dir}"
    except Exception as e:
        return f"Error: {e}"


# 查询 OpenCV 库的大小
print(f"opencv-python size: {get_package_size_improved('opencv-python')}")
print(f"opencv-contrib-python size: {get_package_size_improved('opencv-contrib-python')}")
